import ast
import contextlib
import inspect
import textwrap
from _ast import FunctionDef
from typing import Dict, Union, Tuple, List, Optional, Callable

from llvmlite import ir

from MelodieAllegro.cython_interface.shared import from_buffer_ptr
from MelodieAllegro.global_mgrs import id_generator

double = ir.DoubleType()
integer = ir.IntType(64)
i32 = ir.IntType(32)
customTypes: Dict[str, ir.Type] = {}
string = ir.PointerType(ir.IntType(8))


# 有效的类型标识符：
# 1. int float str
# 2. 用户自定义的类
# 3. 数组List[xx]表达式。

class Modifier:
    def __init__(self, private: bool = False):
        self.private = private

    def __repr__(self):
        return f"{'private' if self.private else 'public'}"


class PMType:
    def __init__(self, is_pointer: bool):
        self._type = None
        self._children: List[PMType] = []
        self.is_pointer = is_pointer

    def get_type(self):
        raise NotImplementedError

    def build_structure(self):
        raise NotImplementedError

    def get_ref_type(self):
        """
        获取在栈上的引用类型。
        对于基础类型，为值
        对于对象, 为指针
        对于列表，为一个引用块，包含指向数组的指针，以及数组的长度信息。

        :return:
        """
        return self.build_structure()

    def get_size(self) -> int:
        raise NotImplementedError

    @property
    def type(self):
        return self._type

    def get_init_value(self):
        raise NotImplementedError

    def get_ast_constant(self) -> ast.Constant:
        """
        当整理代码结构时，需要将所有的变量定义都提到最上端。
        因此，需要生成一个默认值。

        对于数字类型，默认值为0
        对于对象或者字符串类型，默认值为一个空指针

        :return:
        """
        raise NotImplementedError

    def __repr__(self):
        return f"{self.__class__.__name__} {self.build_structure()}"

    def to_py_type(self):
        raise NotImplementedError


class PMHeapType(PMType):
    REFCOUNT = 0
    SIZE = 4

    def inc_ref(self, builder: ir.IRBuilder, auto_pointer_ptr: ir.LoadInstr) -> None:
        from .exteral_functions.builtin_func_types import builtins
        cmp_val = builder.icmp_signed(">", builder.ptrtoint(auto_pointer_ptr, integer), integer(0))
        with builder.if_then(cmp_val):
            # builtins.raise_error(builder, f"Pointer to autopointer ptr of type `{self.__repr__()}` is NULL")
            refcnt_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.REFCOUNT)])
            last_refcount_name = id_generator.new_id('last_refcount')
            new_refcount_name = id_generator.new_id('new_refcount')

            refcnt = builder.load(refcnt_ptr, name=last_refcount_name)
            refcnt_new = builder.add(refcnt, integer(1), name=new_refcount_name)

            builder.store(refcnt_new, refcnt_ptr)

    def dec_ref(self, builder: ir.IRBuilder, auto_pointer_ptr: ir.LoadInstr, extra_info: str = "") -> None:
        from .exteral_functions.builtin_func_types import builtins
        cmp_val = builder.icmp_signed(">", builder.ptrtoint(auto_pointer_ptr, integer), integer(0))
        with builder.if_then(cmp_val):
            refcnt_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.REFCOUNT)])
            last_refcount_name = id_generator.new_id('last_refcount')
            new_refcount_name = id_generator.new_id('new_refcount')
            refcnt = builder.load(refcnt_ptr, name=last_refcount_name)

            # 判断引用计数目前不为0，否则将抛出异常。
            with builder.if_then(builder.icmp_signed("==", refcnt, integer(0))):
                name = self.name if isinstance(self, PMObjectType) else self.__repr__()
                if name == "":
                    raise NotImplementedError
                builtins.raise_error(builder,
                                     f"Refcount for \"{name}\" is already zero when decreasing reference count!\n{extra_info}")
            refcnt_new = builder.sub(refcnt, integer(1), name=new_refcount_name)
            if self.has_heap_child():
                with builder.if_then(builder.icmp_signed("==", refcnt_new, integer(0))):
                    self.build_on_refcnt_zero(builder, auto_pointer_ptr)
            builder.store(refcnt_new, refcnt_ptr)

    def has_heap_child(self):
        raise NotImplementedError

    def build_on_refcnt_zero(self, builder: ir.IRBuilder, obj):
        raise NotImplementedError

    def get_refcnt(self, builder: ir.IRBuilder, auto_pointer_ptr: ir.LoadInstr):
        refcnt_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.REFCOUNT)])
        return builder.load(refcnt_ptr)

    def auto_pointer_type(self) -> ir.LiteralStructType:
        raise NotImplementedError

    def get_ref_type(self) -> ir.PointerType:
        return ir.PointerType(self.auto_pointer_type())

    def get_ptr(self, builder: ir.IRBuilder, arr_ref, name="") -> ir.LoadInstr:
        raise NotImplementedError

    def create(self, builder: ir.IRBuilder, meta, name: str = "") -> ir.LoadInstr:
        pass


class PMBasicType(PMType):
    def __init__(self, initial_value_type, is_pointer=False):
        super(PMBasicType, self).__init__(is_pointer)
        self.is_pointer = is_pointer
        assert self._type in (integer, double, None), self._type
        self._type: ir.Type = initial_value_type

    def build_structure(self):
        return self._type

    def get_ast_constant(self) -> ast.Constant:
        if self._type == double:
            return ast.Constant(value=0.0, kind=None)
        elif self._type == integer:
            return ast.Constant(value=0, kind=None)
        else:
            raise NotImplementedError(f"type {self._type} is not supported")

    def get_size(self) -> int:
        return 8

    def get_init_value(self):
        return ir.Constant(self.type, 0)

    def to_py_type(self):
        if self._type == double:
            return float
        elif self._type == integer:
            return int
        else:
            raise NotImplementedError

    def to_empty_value(self):
        if self._type == double:
            return 0.0
        elif self._type == integer:
            return 0
        else:
            raise NotImplementedError

    def __repr__(self):
        return f"<PMBasicType {str(self._type)}>"


class PMStringType(PMType):
    def __init__(self):
        super().__init__(True)

    def build_structure(self):
        return ir.PointerType(ir.IntType(8))

    def create(self, builder: ir.IRBuilder, byts: bytes):
        from .exteral_functions.builtin_func_types import builtins
        m_ptr = builtins.malloc(builder, [integer(len(byts))])

        typ = ir.ArrayType(ir.IntType(8), len(byts))

        str_ptr = builder.bitcast(m_ptr, typ.as_pointer())

        val = ir.Constant(typ, bytearray(byts))
        builder.store(val, str_ptr)
        return builder.bitcast(m_ptr, string)

    def get_size(self) -> int:
        return 8

    def get_ast_constant(self) -> ast.Constant:
        return ast.Constant(value="", kind=None)


class PMGenericType(PMType):
    def __init__(self, ):
        super(PMGenericType, self).__init__(False)
        self._type: PMType = None

    def build_structure(self):
        if self._type is None:
            raise ValueError('_type is not initialized')
        return self._type.build_structure()

    def get_size(self) -> int:
        return 0

    def get_ast_constant(self) -> ast.Constant:
        return ast.Constant(value="", kind=None)


class PMIntegerSetType(PMHeapType):
    DATA = 1
    WEAKREF = 2
    LENGTH = 4  # 字节数量

    def __init__(self):
        super(PMIntegerSetType, self).__init__(True)

    def get_ref_type(self) -> ir.PointerType:
        return ir.PointerType(self.auto_pointer_type())

    def has_heap_child(self):
        return True

    def build_on_refcnt_zero(self, builder: ir.IRBuilder, obj):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins
        builtins.integer_set_free(builder, self.get_body_ptr(builder, obj))

    def build_structure(self):
        return ir.PointerType(integer)

    def auto_pointer_type(self):
        return ir.LiteralStructType(
            [integer, self.build_structure(), integer, integer, integer]
        )

    def handle_add(self, builder: ir.IRBuilder, value, name: str):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins
        builtins.add_to_integer_set(builder, value, )

    def get_body_ptr(self, builder: ir.IRBuilder, auto_pointer_ptr):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins
        cmp_val = builder.icmp_signed("<=", builder.ptrtoint(auto_pointer_ptr, integer), integer(0))
        with builder.if_then(cmp_val):
            builtins.raise_error(builder,
                                 f"Null pointer error occurred in IntegerSet because object itself is NULL.")
        body_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.DATA)])
        body_ref = builder.load(body_ptr, id_generator.new_id('obj-body-ptr'))
        cmp_val = builder.icmp_signed("<=", builder.ptrtoint(body_ref, integer), integer(0))
        with builder.if_then(cmp_val):
            builtins.raise_error(builder, f"Null pointer error occurred in IntegerSet because object body is NULL!")
        return body_ref

    def create(self, builder: ir.IRBuilder, name: str):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins

        # 分配智能指针的内存
        auto_pointer_ptr_raw = builtins.malloc_auto_ptr(builder, [integer(3 * 8)])
        auto_pointer_ptr = builder.bitcast(auto_pointer_ptr_raw, self.get_ref_type(), name=name)

        # 分配数据部分的内存
        # 数据部分只有一个指针，int64型即可。
        ptr = builtins.new_integer_set(builder)
        pointer = builder.bitcast(ptr, self.build_structure())

        # 初始化智能指针
        refcnt_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.REFCOUNT)])
        builder.store(integer(1), refcnt_ptr)

        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.DATA)])
        builder.store(pointer, pointer_ptr)

        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.WEAKREF)])
        builder.store(integer(0), pointer_ptr)

        # 大小为8个字节
        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.LENGTH)])
        builder.store(integer(1), pointer_ptr)

        return auto_pointer_ptr


class PMArrayType(PMHeapType):
    DATA = 1
    WEAKREF = 2
    LENGTH = 3

    def __init__(self, value_type: Union[double, integer, PMType], is_pointer=False):
        super(PMArrayType, self).__init__(is_pointer)
        if value_type == double or value_type == integer:
            value_type = PMBasicType(value_type)
        assert isinstance(value_type, PMType), value_type
        self._children = [value_type]
        self._type = 0

    def build_structure(self):
        s = self._children[0].get_ref_type()
        arr = ir.PointerType(s)
        return arr

    def has_heap_child(self):
        return isinstance(self._children[0], PMHeapType)

    def auto_pointer_type(self):
        return ir.LiteralStructType(
            [integer, ir.PointerType(self._children[0].get_ref_type()), integer, integer, integer]
        )

    def get_ast_constant(self) -> ast.List:
        return ast.List(elts=[])

    def from_buffer_with_auto(self, builder: ir.IRBuilder, auto_ptr_pointer, name=""):
        auto_pointer_ptr = builder.bitcast(
            builder.inttoptr(
                auto_ptr_pointer,
                self.get_ref_type()
            ),
            self.get_ref_type(),
            name=name)
        return auto_pointer_ptr

    @staticmethod
    def from_buffer(builder: ir.IRBuilder, buffer_ptr_integer, length: ir.LoadInstr, elem_type: str, name: str, ):
        assert elem_type in {'double', 'integer'}
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins
        assert buffer_ptr_integer.type == integer, buffer_ptr_integer

        # 数组类型
        arr_type = PMArrayType(integer if elem_type == 'integer' else 'double')
        # 指向缓冲区的指针
        pointer = builder.inttoptr(buffer_ptr_integer, arr_type.build_structure())

        auto_pointer_ptr_raw = builtins.malloc_auto_ptr(builder, [integer(5 * 8)])
        auto_pointer_ptr = builder.bitcast(auto_pointer_ptr_raw, arr_type.get_ref_type(), name=name)

        # 初始化智能指针
        refcnt_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(PMArrayType.REFCOUNT)])
        builder.store(integer(1), refcnt_ptr)

        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(PMArrayType.DATA)])
        builder.store(pointer, pointer_ptr)

        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(PMArrayType.WEAKREF)])
        builder.store(integer(0), pointer_ptr)

        length_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(PMArrayType.LENGTH)])
        builder.store(length, length_ptr)

        max_length_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(PMArrayType.SIZE)])
        builder.store(builder.mul(length, integer(8)), max_length_ptr)

        return auto_pointer_ptr, arr_type

    def create_by_var(self, builder: ir.IRBuilder, length: ir.LoadInstr, name: str):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins

        # 分配智能指针的内存
        auto_pointer_ptr_raw = builtins.malloc_auto_ptr(builder, [integer(5 * 8)])
        auto_pointer_ptr = builder.bitcast(auto_pointer_ptr_raw, self.get_ref_type(), name=name)

        # 分配数据部分的内存
        bytes_ = builder.mul(integer(self._children[0].get_size()), length)  # self._children[0].get_size() * length
        ptr = builtins.malloc(builder, [bytes_])
        pointer = builder.bitcast(ptr, self.build_structure())

        # 初始化智能指针
        refcnt_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.REFCOUNT)])
        builder.store(integer(1), refcnt_ptr)

        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.DATA)])
        builder.store(pointer, pointer_ptr)

        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.WEAKREF)])
        builder.store(integer(0), pointer_ptr)

        length_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.LENGTH)])
        builder.store(length, length_ptr)

        size_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.SIZE)])
        builder.store(bytes_, size_ptr)
        return auto_pointer_ptr

    def create(self, builder: ir.IRBuilder, length, name: str):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins

        # 分配智能指针的内存
        auto_pointer_ptr_raw = builtins.malloc_auto_ptr(builder, [integer(5 * 8)])
        auto_pointer_ptr = builder.bitcast(auto_pointer_ptr_raw, self.get_ref_type(), name=name)

        # 分配数据部分的内存
        bytes_ = self._children[0].get_size() * length
        ptr = builtins.malloc(builder, [integer(bytes_)])
        pointer = builder.bitcast(ptr, self.build_structure())

        # 初始化智能指针
        refcnt_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.REFCOUNT)])
        builder.store(integer(1), refcnt_ptr)

        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.DATA)])
        builder.store(pointer, pointer_ptr)

        pointer_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.WEAKREF)])
        builder.store(integer(0), pointer_ptr)

        length_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.LENGTH)])
        builder.store(integer(length), length_ptr)

        size_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.SIZE)])
        builder.store(integer(bytes_), size_ptr)
        return auto_pointer_ptr

    def get_size(self) -> int:
        return 16

    def get_length(self, builder: ir.IRBuilder, arr_ref, name=""):
        length = builder.gep(arr_ref, [i32(0), i32(self.LENGTH)])
        return builder.load(length, name)

    def get_max_length(self, builder: ir.IRBuilder, arr_ref, name=""):
        length = builder.gep(arr_ref, [i32(0), i32(self.SIZE)])
        return builder.load(length, name)

    def set_length(self, builder: ir.IRBuilder, arr_ref, new_length):
        length = builder.gep(arr_ref, [i32(0), i32(self.LENGTH)])
        return builder.store(new_length, length)

    def get_ptr(self, builder: ir.IRBuilder, arr_ref, name=""):
        addr = builder.gep(arr_ref, [i32(0), i32(self.DATA)])
        return builder.load(addr, name)

    def reset_ptr(self, builder: ir.IRBuilder, arr_ref, new_data_ptr):
        addr = builder.gep(arr_ref, [i32(0), i32(self.DATA)])
        builder.store(new_data_ptr, addr, )

    @contextlib.contextmanager
    def bound_check(self, builder, arr_ref, index):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins
        cmp_result1 = builder.icmp_signed(
            "<=", integer(0), index)
        length = self.get_length(builder, arr_ref)
        cmp_result2 = builder.icmp_signed(
            "<", index, length
        )
        with builder.if_else(builder.and_(cmp_result2, cmp_result1)) as (then, orelse):
            with then:
                pass
            with orelse:
                builtins.raise_index_error(builder, integer(0), length, index)

    def getitem(self, builder: ir.IRBuilder, arr_ref, index: integer, name=""):
        ptr = self.get_ptr(builder, arr_ref, name)
        self.bound_check(builder, arr_ref, index)
        item_ptr = builder.gep(ptr, [index])
        return builder.load(item_ptr, name), self._children[0]

    def setitem(self, builder: ir.IRBuilder, arr_ref, index: integer, value, name="", bound_check=True):
        """

        :param builder:
        :param arr_ref:
        :param index:
        :param value:
        :param name:
        :return:
        """

        ptr = self.get_ptr(builder, arr_ref, name)
        if bound_check:
            self.bound_check(builder, arr_ref, index)
        item_ptr = builder.gep(ptr, [index])
        builder.store(value, item_ptr)

    def build_on_refcnt_zero(self, builder: ir.IRBuilder, arr):
        if isinstance(self._children[0], PMHeapType):
            last_block = builder.block
            counter_var_ptr = builder.alloca(integer)
            builder.store(integer(0), counter_var_ptr)
            counter_var = builder.load(counter_var_ptr, id_generator.new_id('counter'))
            rm_loop_block = builder.append_basic_block(id_generator.new_id("rm-loop"))
            after_rm_loop_block = builder.append_basic_block(id_generator.new_id("after-rm-loop"))
            builder.branch(rm_loop_block)

            builder.position_at_end(rm_loop_block)
            counter_var_phi = builder.phi(integer)
            counter_var_phi.add_incoming(counter_var, last_block)
            counter_var = counter_var_phi

            item_value, item_type = self.getitem(builder, arr, counter_var)
            assert isinstance(item_type, PMHeapType)
            # If not null, decrease ref count!
            with builder.if_then(builder.icmp_signed(">=", builder.ptrtoint(item_value, integer), integer(0))):
                item_type.dec_ref(builder, item_value, "When parent array refcount is zero")

            with builder.if_else(
                    builder.icmp_signed(">=", counter_var,
                                        builder.sub(self.get_length(builder, arr), integer(1))
                                        )
            ) as (then, orelse):
                with then:
                    builder.branch(after_rm_loop_block)
                with orelse:
                    counter_var = builder.add(counter_var, integer(1))
                    builder.branch(rm_loop_block)
                    counter_var_phi.add_incoming(counter_var, builder.block)
            # counter_var_phi.add_incoming(counter_var, builder.block)
            builder.branch(after_rm_loop_block)
            builder.position_at_end(after_rm_loop_block)


class PMStructType(PMType):
    def __init__(self, value_types: Dict[str, Union[PMType, double, integer]], is_pointer=False):
        super(PMStructType, self).__init__(is_pointer)
        self.value_types_mapping = {}
        for i, value_type in enumerate(value_types.items()):
            k, typ = value_type
            if typ == double or typ == integer:
                typ = PMBasicType(typ)
            self._children.append(typ)
            self.value_types_mapping[k] = i

    def initialize(self, builder: ir.IRBuilder, ptr):
        typ = ir.LiteralStructType([child.build_structure() for child in self._children])
        c = ir.Constant(typ, [child.get_init_value() for child in self._children])
        builder.store(c, ptr)

    def build_structure(self):
        st = ir.LiteralStructType([child.build_structure() for child in self._children])
        if self.is_pointer:
            st = ir.PointerType(st)
        return st

    def get_size(self) -> int:
        s = 0
        for child in self._children:
            s += child.get_size()
        return s

    def create(self, builder: ir.IRBuilder, name: str):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins
        if self.is_pointer:
            ptr = builtins.malloc(builder, [integer(self.get_size())])
            pointer = builder.bitcast(ptr, self.build_structure())
            return pointer
        else:
            type = self.build_structure()
            ptr = builder.alloca(type, name=name + '_ptr')
            self.initialize(builder, ptr)
            return builder.load(ptr, name=name)

    def getattr(self, builder: ir.IRBuilder, struct, attr_name: str, name: str):
        offset: int = self.value_types_mapping[attr_name]
        if self.is_pointer:
            ptr = struct
            value_p = builder.gep(ptr, [i32(0), i32(offset)])
            return builder.load(value_p, name)
        else:
            val = builder.extract_value(struct, offset, name)
            return val

    def setattr(self, builder: ir.IRBuilder, struct, attr_name: str, value, name: str = ""):
        offset: int = self.value_types_mapping[attr_name]
        if self.is_pointer:
            ptr = struct
            value_p = builder.gep(ptr, [i32(0), i32(offset)])
            builder.store(value, value_p)
        else:
            val = builder.insert_value(struct, value, offset, name=name)
            return val


class PMFunctionType(PMType):
    def __init__(self, name: str, ret_type: PMType, arg_types: List[PMType],
                 is_method=False,
                 owner: "PMObjectType" = None,
                 modifier: "Modifier" = None):
        super().__init__(False)
        self.name = name
        assert (isinstance(ret_type, PMType) or ret_type is None)
        self.ret_type = ret_type
        self.arg_types = arg_types
        self.is_method = is_method
        self.ir_func: Optional[ir.Function] = None

        self.owner: Optional[PMObjectType] = owner
        self.modifier: Optional[Modifier] = modifier

        self.ast: Optional[FunctionDef] = None
        self.locals: Dict[str, PMType] = {}

    def build_structure(self):
        return self.ir_func

    def func_name_in_ir(self):
        """
        在IR中显示的名称
        比如一个类的方法，可能互相是相同的。
        :return:
        """
        return ('.method.' + self.owner.name + '.' + self.name) if self.is_method else self.name

    def create(self, module: ir.Module, func: Callable = None,
               global_functions: Optional[Dict[str, 'PMFunctionType']] = None,
               classes: Optional[Dict[str, 'PMObjectType']] = None
               ):
        """
        将此函数生成IR

        :param module:
        :param func:
        :param global_functions:
        :param classes:
        :return:
        """
        if func is None:
            func = self.ast
        else:
            func = ast.parse(textwrap.dedent(inspect.getsource(func))).body[0] if callable(func,
                                                                                           ) else func
        assert isinstance(func, FunctionDef)

        from .ast_parser import FunctionVisitor
        types = self.locals.copy()
        if self.is_method:
            types['self'] = self.owner

        fv = FunctionVisitor(module, types=types, func_name_in_ir=self.func_name_in_ir())
        global_functions = {} if global_functions is None else global_functions
        classes = {} if classes is None else classes
        fv.functions = global_functions.copy()
        fv.custom_types = classes.copy()
        self.ir_func = fv.start_visit(func, self.ret_type)

    def call_func(self, builder: ir.IRBuilder, args: List, name=""):
        assert self.ir_func is not None, "This function may not be scanned yet!"
        return builder.call(self.ir_func, args, name=name)

    def call_func_with_ir(self, builder: ir.IRBuilder, ir_func: ir.Function, args: [], name=""):
        assert ir_func is not None
        return builder.call(ir_func, args, name=name)

    def __repr__(self):
        return f"<PMFuncdef ismethod:{self.is_method} {'owner: ' + str(self.owner) if self.is_method else ''} > "


class PMObjectType(PMHeapType):
    """
    All objects are referenced by pointer.

    所有对象类型都是由指针来索引的。
    """
    PROP = 0
    METHOD = 1

    # The first slot of auto pointer stores reference count,
    # the following slots are listed below.
    # 智能指针的引用计数为第一位(0)，其后的各位如下所示
    BODY = 1
    WEAKREF = 2

    AUTO_POINTER_SIZE = 5
    SIZE = 4

    def __init__(self, name: str, prop_types: Dict[str, Tuple[Union[PMType, double, integer], Modifier]],
                 parents: List['PMObjectType'] = None, methods: List[PMFunctionType] = None):
        super().__init__(True)
        self.name = name

        # Properties of this object.
        self.props: List[Tuple[PMType, Modifier, str]] = []

        # Map from property name to offset.
        self.prop_types_mapping: Dict[str, int] = {}

        # Parent objects, supporting multiple inheritance.
        self.parents: List[PMObjectType] = parents if parents is not None else []
        assert isinstance(self.parents, list)

        # The offset of each parent classes
        self.inherit_offset: Dict[str, int] = {}

        # All methods, including defined in this class or inherited from parent
        self.methods: Dict[str, PMFunctionType] = {}

        # Methods that defined in this class.
        child_methods: Dict[str, PMFunctionType] = {f.name: f for f in methods} if methods is not None else {}
        self.child_methods = child_methods

        # Initialize the inheritance offset of properties and
        # inherit the methods from parent.
        pointer = 0
        for parent in self.parents:
            self.inherit_offset[parent.name] = pointer
            self.props.extend([prop for prop in parent.props])
            self.methods.update(parent.methods)
            pointer += 1
        self.inherit_offset[self.name] = 0

        # Update the methods map of children.
        # If method of child has the same name as its parent's, the parent method
        # will be overridden.
        self.methods.update(child_methods)

        # Append the props defined in children
        # to the table of properties.
        for k, prop in prop_types.items():
            typ, modifier = prop
            if typ == double or typ == integer:
                typ = PMBasicType(typ)
            prop = (typ, modifier, k)
            self.props.append(prop)

        # Reverse the props list to create the property map
        # { PropertyName } --->  { Index }
        i = len(self.props) - 1
        for prop in reversed(self.props):
            name = prop[2]
            if name not in self.prop_types_mapping:
                self.prop_types_mapping[name] = i
            i -= 1

        self.on_ref_zero: Optional[ir.Function] = None

    def get_property_type(self, prop_name) -> PMType:
        offset = self.prop_types_mapping[prop_name]
        return self.props[offset][0]

    @staticmethod
    def same_type(ot1: "PMObjectType", ot2: "PMObjectType"):
        """
        Tell if object type `ot1` and object type `ot2` are of the same type.

        :param ot1:
        :param ot2:
        :return:
        """
        return ot1.build_structure().pointee == ot2.build_structure().pointee

    def call_method_by_functype(self, builder: ir.IRBuilder, method: PMFunctionType, obj_type: "PMObjectType",
                                obj_self,
                                args_except_self: List = None, name=""):
        args_except_self = [] if args_except_self is None else args_except_self

        assert isinstance(obj_type, PMObjectType), obj_type
        if PMObjectType.same_type(obj_type, self):
            obj_type.inc_ref(builder, obj_self)
            return method.call_func(builder, [obj_self] + args_except_self, name)
        else:
            bitcasted = obj_type.bitcast_to_parent(builder, self, obj_self)
            obj_type.inc_ref(builder, obj_self)
            return method.call_func(builder, [bitcasted] + args_except_self, name)

    def call_method_from_class(self, builder: ir.IRBuilder, method_name: str, obj_type: "PMObjectType", obj_self,
                               args_except_self: List = None, name=""):
        return self.call_method(builder, method_name, obj_type, obj_self, args_except_self, name), self.get_method(
            method_name).ret_type

    def call_method(self, builder: ir.IRBuilder, method_name: str, obj_type: "PMObjectType", obj_self,
                    args_except_self: List = None, name=""):
        method: PMFunctionType = self.get_method(method_name)
        return self.call_method_by_functype(builder, method, obj_type, obj_self, args_except_self, name)

    def get_method(self, method_name: str) -> PMFunctionType:
        if method_name in self.child_methods:
            return self.child_methods[method_name]
        elif method_name in self.methods:
            return self.methods[method_name]
        else:
            raise KeyError(method_name)

    def bitcast_to_parent(self, builder: ir.IRBuilder, parent: "PMObjectType", obj):
        """
        Bitcast object type to its parent

        WARNING: This step will create a new auto pointer by weak-reference.

        :param builder:
        :param parent:
        :param obj: the auto-pointer of an object
        :return:
        """
        body_ptr = self.get_body_ptr(builder, obj)
        bitcasted_body = builder.bitcast(body_ptr, parent.build_structure(),
                                         id_generator.new_id(f'convert{self.name}-to-{parent.name}'))
        tmp_bitcast_name = id_generator.new_id('temp-bitcast')
        return self.create_auto_pointer(builder, bitcasted_body, name=tmp_bitcast_name, obj=parent, weakref=True)

    def print_memory_layout(self):
        print("=============Memory space of Object Type is================")
        for i, item in enumerate(self.props):
            print(i, item[1], item[0], item[2])

    def get_size(self) -> int:
        s = 0
        for child, _, _ in self.props:
            s += child.get_size()
        return s

    def build_structure(self):
        st = ir.LiteralStructType([typ.get_ref_type() for typ, _, _ in self.props])
        return ir.PointerType(st)

    def create_type(self, module: ir.Module, functions=None, classes=None):
        """
        创建此类型的所有方法

        :param module:
        :param functions:
        :param classes:
        :return:
        """
        self.on_ref_zero = self.build_on_refcnt_zero_function(module)
        methods = self.child_methods.copy()
        if '__init__' in methods:
            methods.pop('__init__').create(module, global_functions=functions, classes=classes)

        for name, method in methods.items():
            method.create(module, global_functions=functions, classes=classes)

    def auto_pointer_type(self) -> ir.LiteralStructType:
        return ir.LiteralStructType([
            integer,
            self.build_structure(),
            integer,
            integer,
            integer
        ])

    def copy_props(self, builder: ir.IRBuilder, obj, data_ptr):
        """
        Copy property from
        :return:
        """
        obj_data_ptr = self.get_ptr(builder, obj, )

    def write_to_python_object(self, buffer_ptr: int, py_obj: object, ignore_object=True):
        """
        将内存中的对象数据，写入python对象中。

        :param props:
        :return:
        """
        props_to_get: List[Tuple] = []  # [(name, type,  offset)]
        for prop_type, _, prop_name in self.props:
            offset = self.prop_types_mapping[prop_name]
            props_to_get.append((prop_name, prop_type, offset))

        from_buffer_ptr(buffer_ptr, props_to_get, py_obj, ignore_object)

    def create_in_memory(self, props: Dict[str, Union[int, float]]):
        from MelodieAllegro.cython_interface.shared import write_to_buffer, alloc_obj_memory
        memory_length = len(self.props)
        order: List[str] = []
        value_dict: Dict[str, Union[int, float]] = {}
        for prop in self.props:
            pm_type, modifier, prop_name = prop
            order.append(prop_name)
            if prop_name in props:
                if isinstance(pm_type, PMBasicType):
                    assert isinstance(props[prop_name],
                                      pm_type.to_py_type()), f"Property name {prop_name} value {props[prop_name]} " \
                                                             f"is not the same as required type {pm_type.to_py_type()}"
                    value_dict[prop_name] = props[prop_name]
                elif isinstance(pm_type, PMObjectType):
                    assert isinstance(props[prop_name], int), props[prop_name]
                    value_dict[prop_name] = props[prop_name]
                elif isinstance(pm_type, PMArrayType):
                    value_dict[prop_name] = props[prop_name]
                else:
                    raise NotImplementedError(pm_type)

            else:
                if isinstance(pm_type, PMBasicType):
                    value_dict[prop_name] = pm_type.to_empty_value()
                else:
                    value_dict[prop_name] = 0

        return alloc_obj_memory([value_dict], memory_length, order)

    def from_buffer(self, builder: ir.IRBuilder, auto_ptr_pointer, name: str = ""):
        auto_pointer_ptr = builder.bitcast(builder.inttoptr(auto_ptr_pointer, self.get_ref_type()), self.get_ref_type(),
                                           name=name)
        return auto_pointer_ptr

    def get_autoptr_size(self, builder):
        from MelodieAllegro.macros import PMMacroGenerator
        return PMMacroGenerator(builder).get_size(self.auto_pointer_type())

    def create_auto_pointer(self, builder: ir.IRBuilder, body_pointer, name: str = "",
                            obj: 'PMObjectType' = None, weakref: bool = False):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins
        if obj is None:
            obj = self
        # 分配智能指针的内存

        auto_pointer_ptr_raw = builtins.malloc_auto_ptr(builder, [self.get_autoptr_size(builder)])
        auto_pointer_ptr = builder.bitcast(auto_pointer_ptr_raw, obj.get_ref_type(), name=name)

        refcnt_addr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.REFCOUNT)])
        builder.store(integer(1), refcnt_addr)

        body_addr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.BODY)])
        builder.store(body_pointer, body_addr)

        weakref_addr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.WEAKREF)])
        builder.store(integer(int(weakref)), weakref_addr)

        size_addr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.SIZE)])
        assert self.get_size() % 8 == 0, self.get_size()
        builder.store(integer(int(self.get_size() / 8)), size_addr)

        return auto_pointer_ptr

    def create(self, builder: ir.IRBuilder, args: List, name: str = ""):
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins

        # Allocate memory for properties
        structure_type = self.build_structure()
        ptr = builtins.malloc(builder, [integer(self.get_size())])
        body_pointer = builder.bitcast(ptr, structure_type, name=name + '-data')

        # Create the auto pointer, and here it is Strong reference.
        auto_pointer_ptr = self.create_auto_pointer(builder, body_pointer, name, weakref=False)

        if self.methods.get('__init__') is not None:
            self.get_method('__init__').owner.call_method(builder, '__init__', self, auto_pointer_ptr, args)
        return auto_pointer_ptr

    def tell_attr_type(self, attr_name: str):
        """
        Tell if one attribute is method or property.

        :param attr_name:
        :return:
        """
        if attr_name in self.prop_types_mapping.keys():
            return self.PROP
        else:
            assert attr_name in self.methods, (attr_name, self.methods)
            return self.METHOD

    def get_attr_type(self, attr_name) -> PMType:
        """
        Get the type of one attribute by name.

        :param attr_name:
        :return:
        """
        offset: int = self.prop_types_mapping[attr_name]
        typ = self.props[offset][0]
        assert isinstance(typ, PMType)
        return typ

    def get_body_ptr(self, builder: ir.IRBuilder, auto_pointer_ptr: ir.LoadInstr, meta=""):
        """
        Get the pointer to properties data in heap.

        :param builder:
        :param auto_pointer_ptr:
        :return:
        """
        from MelodieAllegro.exteral_functions.builtin_func_types import builtins
        cmp_val = builder.icmp_signed("<=", builder.ptrtoint(auto_pointer_ptr, integer), integer(0))
        with builder.if_then(cmp_val):
            builtins.raise_error(builder,
                                 f"Null pointer error occurred when {meta} because object itself is NULL.")
        body_ptr = builder.gep(auto_pointer_ptr, [i32(0), i32(self.BODY)])
        body_ref = builder.load(body_ptr, id_generator.new_id('obj-body-ptr'))
        cmp_val = builder.icmp_signed("<=", builder.ptrtoint(body_ref, integer), integer(0))
        with builder.if_then(cmp_val):
            builtins.raise_error(builder, f"Null pointer error occurred when {meta} because object body is NULL!")
        return body_ref

    def setattr_by_offset(self, builder: ir.IRBuilder, obj, attr_offset, value):
        ptr = self.get_body_ptr(builder, obj, meta=f'setattr_by_offset{attr_offset}')
        value_p = builder.gep(ptr, [i32(0), i32(attr_offset)])
        builder.store(value, value_p)

    def getattr_by_offset(self, builder: ir.IRBuilder, obj, attr_offset: int, name: str):
        """
        Get the attribute.

        :return:
        """
        ptr = self.get_body_ptr(builder, obj, meta=f'getattr_by_offset, {attr_offset}')
        value_p = builder.gep(ptr, [i32(0), i32(attr_offset)])
        attr_type = self.props[attr_offset][0]
        return builder.load(value_p, name), attr_type

    def getattr(self, builder: ir.IRBuilder, obj, attr_name: str, name: str):
        """
        Get the attribute.

        :return:
        """
        offset: int = self.prop_types_mapping[attr_name]
        ptr = self.get_body_ptr(builder, obj, meta=f"getting attribute '{attr_name}' from `{self.name}`")
        value_p = builder.gep(ptr, [i32(0), i32(offset)])
        return builder.load(value_p, name)

    def attr_address(self, builder: ir.IRBuilder, obj, attr_name: str, name: str):
        offset: int = self.prop_types_mapping[attr_name]
        ptr = self.get_body_ptr(builder, obj, meta=f"getting address of attribute '{attr_name}' from `{self.name}`")
        value_p = builder.gep(ptr, [i32(0), i32(offset)])
        return value_p

    def setattr(self, builder: ir.IRBuilder, obj, attr_name: str, value, name: str = ""):
        """
        Set the attribute

        :return:
        """
        offset: int = self.prop_types_mapping[attr_name]
        ptr = self.get_body_ptr(builder, obj, meta=f"setting attribute '{attr_name}' to `{self.name}`")
        value_p = builder.gep(ptr, [i32(0), i32(offset)])
        builder.store(value, value_p)

    def setitem(self, builder: ir.IRBuilder, obj, key, value):
        getitem_method = self.get_method('__setitem__')
        parent = getitem_method.owner
        self.inc_ref(builder, obj)
        if not PMObjectType.same_type(parent, self):
            bitcasted = self.bitcast_to_parent(builder, parent, obj)
            args = [bitcasted, key, value]
        else:
            args = [obj, key, value]
        getitem_method.call_func(builder, args,
                                 id_generator.new_id('obj-setitem')
                                 )

    def getitem(self, builder: ir.IRBuilder, obj, index, name=""):
        """
        Getitem by calling the magic function __getitem__

        :return:
        """
        getitem_method = self.get_method('__getitem__')
        parent = getitem_method.owner
        self.inc_ref(builder, obj)
        if not PMObjectType.same_type(parent, self):
            bitcasted = self.bitcast_to_parent(builder, parent, obj)
            args = [bitcasted, index]
        else:
            args = [obj, index]
        return getitem_method.call_func(builder, args, name), getitem_method.ret_type

    def __repr__(self):
        return f"<PMObjectType {self.name} {self.build_structure()}>"

    def has_heap_child(self):
        for prop in self.props:
            if isinstance(prop[0], PMHeapType):
                return True
        return False

    def build_on_refcnt_zero(self, builder: ir.IRBuilder, obj):
        builder.call(self.on_ref_zero, [obj], 'on-ref-zero')

    def build_on_refcnt_zero_function(self, module: ir.Module):
        func_ty = ir.FunctionType(ir.VoidType(), [self.get_ref_type()])
        func = ir.Function(module, func_ty, name=f"on-{self.name}-refcnt-zero")
        builder = ir.IRBuilder(func.append_basic_block('entry'))
        obj_to_release = func.args[0]
        for prop in self.props:
            prop_type = prop[0]
            prop_name = prop[2]
            assert isinstance(prop_type, PMType)
            if isinstance(prop_type, PMHeapType):
                prop_type.dec_ref(builder, self.getattr(builder, obj_to_release, prop_name, 'getattr-for-ref'),
                                  f"for attribute {obj_to_release}.{prop_name}")
        builder.ret_void()
        return func

    def has_attr(self, attr_name: str):
        return attr_name in self.prop_types_mapping


def pm_function(f):
    def wrapper(*args, **kwargs):
        return f(*args, **kwargs)

    return wrapper
